import os
import csv
import cv2
import numpy as np
from pyquaternion import Quaternion
from computeAngle import *


def extractFrame(fileName, sample_rate=4, dsize=(160, 90)):
    """
    extract video frame

    :param fileName: the name of extracted video
    :param sample_rate: the interval of extracting frame
    :param dsize: the frame size of extracting
    :return: None (save timeStamp.txt)
    """

    videoPath = 'D:/VR_project/LiveDeep_All/videos/' + fileName
    savePath = 'D:/VR_project/ViewPrediction/frames/' + fileName.split('.')[0]
    if not os.path.exists(savePath):
        os.makedirs(savePath)
    else:
        for i in os.listdir(savePath):
            os.remove(savePath + '/' + i)

    timeList = []

    cameraCapture = cv2.VideoCapture(videoPath)
    fps = cameraCapture.get(cv2.CAP_PROP_FPS)
    iVideoTime = cameraCapture.get(cv2.CAP_PROP_FRAME_COUNT) / cameraCapture.get(cv2.CAP_PROP_FPS)
    frames = cameraCapture.get(cv2.CAP_PROP_FRAME_COUNT)

    frameId, idx = 1, 1

    while idx <= frames:
        success, frame = cameraCapture.read()
        if not success:
            print(f'\nRead finish! Have gotten {frameId-1} sampled frames!')
            break

        milliseconds = cameraCapture.get(cv2.CAP_PROP_POS_MSEC)
        seconds = milliseconds // 1000
        milliseconds = milliseconds % 1000

        if idx % int(fps/sample_rate) == 0:
            cv2.imwrite(savePath+f'/{frameId}.jpg', cv2.resize(frame, dsize))
            frameId += 1
            print('\r{0}.{1}/{2}'.format(int(seconds), int(milliseconds), round(iVideoTime,3)), end='')
            timeList.append(float(str(int(seconds)) + '.' + str(int(milliseconds))))

        idx += 1
    np.savetxt(timeFoVPath + fileName.split('.')[0]+'.txt', np.array(timeList), fmt='%.3f')

    cv2.destroyAllWindows()
    cameraCapture.release()


def createFixationMap(userId, videoId, timeList):
    """get and save fixationMap by extracting FoVData

    :param userId: user Id
    :param videoId: the extracted video's Id
    :param timeList: sample frame's timeStamp
    :return: None (save fixationMap.npy)
    """

    # part1: get FoVData according to userId and videoId
    FoVData = []
    userFile = f'D:/VR_project/LiveDeep_All/vr-dataset/Experiment_1/{userId}/video_{videoId}.csv'
    if not os.path.exists(userFile):
        exit(f"{userFile} not exists")

    with open(userFile) as csvFile:
        csv_reader = csv.reader(csvFile)
        next(csv_reader)
        i, t_len = 0, len(timeList)
        for row in csv_reader:
            if float(row[1]) >= timeList[i]:
                v0 = [0, 0, 1]
                q = Quaternion([float(row[5]), -float(row[4]), float(row[3]), -float(row[2])])
                FoVData.append(q.rotate(v0))
                i += 1
                if i == t_len:
                    break

    # part2: get fixationMap and save
    fixationMap = []
    for fov in FoVData:
        theta, phi = vector_to_ang(fov)
        hi, wi = ang_to_geoxy(theta, phi, H, W)
        result = np.zeros(shape=(H, W))
        result[H - hi - 1, W - wi - 1] = 1
        fixationMap.append(result)
    FoVPath = f'D:/VR_project/ViewPrediction/frames/timeStampFoV/{videoName[videoId]}'
    if not os.path.exists(FoVPath):
        os.makedirs(FoVPath)
    np.save(FoVPath + f'/{userId}.npy', np.array(fixationMap))
    print(f'save {videoName[videoId]}\'s fixationMap for user={userId}')



if __name__ == "__main__":
    fileList = ['1-1-Conan Gore Fly.mp4', '1-2-Front.mp4', '1-3-360 Google Spotlight Stories_ HELP.mp4',
                 '1-4-Conan Weird Al.mp4', '1-5-TahitiSurf.mp4', '1-6-Falluja.mp4',
                '1-7-Cooking Battle.mp4', '1-8-Football.mp4', '1-9-Rhinos.mp4']
    videoName = ['Conan1', 'Skiing', 'Alien', 'Conan2', 'Surfing', 'War', 'Cooking', 'Football', 'Rhinos']

    timeFoVPath = 'D:/VR_project/ViewPrediction/frames/timeStampFoV/'

    videoId = 6
    fileName = fileList[videoId]
    H, W = 90, 160
    # extractFrame(fileName)

    timeList = np.loadtxt(timeFoVPath + fileName.split('.')[0] + '.txt')
    for userId in range(1, 49):
        createFixationMap(userId, videoId, timeList)
